import torch
import torch.nn as nn
import torch.optim as optim
from torch.nn.utils import prune

# 定义一个简单的神经网络模型
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc1 = nn.Linear(784, 512)
        self.fc2 = nn.Linear(512, 10)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 初始化模型、损失函数和优化器
model = SimpleModel()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 训练模型（此处省略）

# 应用剪枝
prune.l1_unstructured(model.fc1, 'weight', amount=0.5)  # 移除50%的参数

# 微调剪枝后的模型（此处省略）


# 定义“教师”模型和“学生”模型
class TeacherModel(nn.Module):
    # ...（模型定义）

class StudentModel(nn.Module):
    # ...（较小的模型定义）

teacher = TeacherModel()
student = StudentModel()

# 加载预训练的“教师”模型权重
teacher.load_state_dict(torch.load('teacher_model.pth'))
teacher.eval()

# 准备数据加载器
train_loader = ...  # 数据加载器定义

# 定义损失函数和优化器
criterion = nn.KLDivLoss(reduction='batchmean')
optimizer = optim.Adam(student.parameters(), lr=0.001)

# 训练“学生”模型
for epoch in range(num_epochs):
    for data, target in train_loader:
        optimizer.zero_grad()

        # 使用“教师”模型生成软标签
        with torch.no_grad():
            teacher_output = teacher(data)
            soft_targets = torch.softmax(teacher_output, dim=1)

        # “学生”模型前向传播
        student_output = student(data)

        # 计算损失
        loss = criterion(torch.log_softmax(student_output, dim=1), soft_targets)

        # 反向传播和优化
        loss.backward()
        optimizer.step()

    print(f'Epoch {epoch+1}, Loss: {loss.item()}')

# 保存训练好的“学生”模型
torch.save(student.state_dict(), 'student_model.pth')

